args <- commandArgs(T)

# "'Rscript
# /home/hsmy/code/ik/TCR/singlecell/XCRAna/TCRana1.R
# /home/hsmy/mob/TCR.list
# /home/hsmy/mob/p.TCRinfo1.xls
# 1'"

# TCRDirListi <- "/home/hsmy/mob/TCR.list"
# outTCRFile <- "/home/hsmy/mob/p.TCRinfo1.xls"
# flag <- 1  #1,all; 2:1A1B,1A2B,2A1B 3:only 1A0B or 0A1B

TCRDirListi <- args[1]
outTCRFile <- args[2]
flag <- args[3]  #1,all; 2:1A1B,1A2B,2A1B 3:only 1A0B or 0A1B

#output two files:
#TCRTable and TCRTable.StandTCR.Bar


#TCRDirListi:
#P4_Ptumor_	/data2/users/dongxingjun/Human_SC_AnaData/P4/Ptumor/TCR/outs

#finish in 2021.11.10
#from cellrange's TCR analysis data get TCR Table

readList <- function(LTable){
  a <- read.table(LTable,sep = "\t")
  sampleList  <- as.vector(a$V1)
  b <- as.vector(a$V2)
  names(b) <- as.vector(a$V1)
  return(b)
}

subnull <- function(rawvec) {
  subs <- c()
  for(i in rawvec)
  {
    if(i == "")
    {
      subs <- c(subs,"None")
    }else
    {
      subs <- c(subs,i)
    }
  }
  return(subs)
}


readTCR <- function(data.dir=NULL,CellType=NULL) {
  #get all TCR info
  contig.loc <- paste0(data.dir,"/filtered_contig_annotations.csv")
  contig <- read.csv(contig.loc)
  contig$d_gene <- subnull(contig$d_gene)
  
  all.cell <- as.vector(contig[,1])
  
  contig[,1] <- all.cell

  #filter cdr3 not none
  contig <- contig[which(!contig$cdr3 =="None"),]

  productiveA.cell <- unique(all.cell[which(tolower(contig$high_confidence)=="true" & contig$chain=="TRA")])
  productiveB.cell <- unique(all.cell[which(tolower(contig$high_confidence)=="true" & contig$chain=="TRB")])
  productiveAB.cell <- union(productiveA.cell,productiveB.cell)

  id <- apply(contig[,c("v_gene","d_gene","j_gene","cdr3")],1, function(x)paste(x,collapse = ":"))
  clonotype2cell <- cbind(id,contig[,c("barcode","chain","raw_clonotype_id")])[which(tolower(contig$high_confidence) == "true"),]
  clonotype2cell <- clonotype2cell[which(clonotype2cell$barcode %in% productiveAB.cell),]

  #get final table
  all.TCR.srt <- lapply(unique(clonotype2cell$barcode), function(x)
  {
    idx <- which(clonotype2cell$barcode==x)
    TCR <- clonotype2cell[idx,]
    
    TRA <- TCR[which(TCR$chain=="TRA"), "id"]
    TRA.count <- length(TRA)
    TRA <- paste(TRA, collapse = "|")
    
    TRB <- TCR[which(TCR$chain=="TRB"), "id"]
    TRB.count <- length(TRB)
    TRB <- paste(TRB, collapse = "|")
    
    TCR.id <- paste(TRA,TRB,sep=";")
    Tcount <- TRA.count+TRB.count
    
    TCR.stat <- paste(TRA.count,"A",TRB.count,"B",sep="")
    
    c(TCR.id,TRA.count,TRB.count,Tcount,TCR.stat)
    
  })

  all.TCR.srt <- do.call(rbind,all.TCR.srt)
  barcode <- unique(clonotype2cell$barcode)
  barcode <- paste0(CellType, barcode)
  rownames(all.TCR.srt) <- barcode
  
  finalTable <- as.data.frame(all.TCR.srt)
  finalTable <- cbind(barcode,finalTable)
  TableColName <- c("barcode","TCRid","ACount","BCount","TotCount","TCRstat")
  colnames(finalTable) <- TableColName

  return(finalTable)
}

#main
a <- readList(TCRDirListi)
cellTypeList <- names(a)

#for every sample
finalTable <- NULL
for(i in seq(1,length(cellTypeList),1)) {
  celltype <- cellTypeList[i]
  dataDir <- as.character(a[i])
  print(celltype)
  print(dataDir)
  tempTable <- readTCR(data.dir=dataDir,CellType=celltype)      
  finalTable <- rbind(finalTable,tempTable)
}

#filter
if(flag == 1)
{
  #all kinds of TCR (XAXB)
  print("all kinds of TCR barcode:")
  finalTable <- finalTable
}else if(flag == 2)
{
  print("TCR:1A1B,2A1B,1A2B,2A2B:")
  #1A1B,2A1B,1A2B
  finalTable <- finalTable[which(finalTable$ACount %in% c(1,2) & finalTable$BCount %in% c(1,2) & finalTable$TotCount %in% c(2,3,4)),]
}else if(flag == 3)
{
  print("TCR:1A0B,0A1B:")
  #only 1A0B or 0A1B
  finalTable <- finalTable[which(finalTable$ACount %in% c(0,1) & finalTable$BCount %in% c(0,1) & finalTable$TotCount %in% c(1)),]
}


#output final table
write.table(finalTable,file = outTCRFile,row.names = F,col.names = T,sep="\t",quote=F)
outBar <- paste(outTCRFile,".StandTCR.Bar",sep="")
print(outBar)
outBarTCRStat <- c("1A1B","1A2B","2A1B","2A2B")
writeLines(as.vector(finalTable[which(finalTable$TCRstat %in% outBarTCRStat),]$barcode),con=outBar,sep="\n")

